package com.timeriver

import com.alibaba.fastjson.{JSON, JSONArray, JSONException, JSONObject}
import org.apache.spark.ml.PipelineModel
import org.apache.spark.ml.feature.LabeledPoint
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.streaming.StreamingQuery
import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}

import scala.collection.mutable.ArrayBuffer

object SparkStreamPredict {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .appName("kafka流数据模型预测")
      .master("local[6]")
      .getOrCreate()

    session.conf.set("spark.sql.streaming.checkpointLocation", "./checkpoint")

    import session.implicits._

    /** 读取 kafka 数据转为 DataFrame */
    val dataFrame: DataFrame = session.readStream
      .format("kafka")
      .option("kafka.bootstrap.servers", "10.0.24.197:9092")
      .option("subscribe", "example-001")
      .option("startingOffsets", "latest")
      .load()

    /** 预测数据 */
    val frame: DataFrame = dataFrame.selectExpr("CAST(value as STRING)")

    val value: Dataset[LabeledPoint] = frame.map(row => {
      /** 可变数组 */
      var arr = ArrayBuffer[Double]()
      var label = 0
      try {
        val jSONObject: JSONObject = JSON.parseObject(row.getString(0))
        label = jSONObject.getInteger("label")
        val array: JSONArray = jSONObject.getJSONArray("data")
        for (i <- 0 until array.size()) {
          arr += array.getJSONObject(i).getDouble("value")
        }
      } catch  {
        case e: JSONException => println(row.getString(0))
      }

      LabeledPoint(
        label, Vectors.dense(arr.toArray)
      )
    })

    val model: PipelineModel = PipelineModel.load("./model/logistic_reg")

    val predict: DataFrame = model.transform(value)
      .select("label", "prediction")


    /** 将结果数据以JSON格式写入kafka中 */
    val query: StreamingQuery = predict.toJSON
      .writeStream
      .format("kafka")
      .option("kafka.bootstrap.servers", "10.0.24.197:9092")
      .option("topic", "t_prediction")
      .start()

    /** 流式程序挂起 */
    query.awaitTermination()
  }
}
